import sys
sys.path.append("/data/python_project/qiufengfeng/nlp_tools")
from nlp_tools.tasks.labeling import ABCLabelingModel
from nlp_tools.tasks.abs_task_model import ABCTaskModel
import re

import sys
if 'win' in sys.platform:
    save_path = r'E:\model_output\ner'
else:
    save_path = r'/home/qiufengfeng/nlp/train_models/ner'
loaded_model = ABCTaskModel.load_model(save_path)
print(loaded_model.get_input_and_output_names_from_model())
def cut_text(text,length):
    textArr = re.findall('.{' + str(length) + "}",text)
    textArr.append(text[(len(textArr) * length) :])
    return textArr

def extract_labels(text,ners):
    ner_reg_list = []
    if ners :
        new_ners = []
        for ner in ners:
            new_ners += ner

        for word,tag in zip([char for char in text],new_ners):
            if tag != "O":
                ner_reg_list.append((word,tag))

    # 输出模型的NER识别结果
    labels = {}
    if ner_reg_list:
        for i,item in enumerate(ner_reg_list):
            if item[1].startswith('B'):
                label = ""
                end = i + 1
                while end <= len(ner_reg_list) - 1 and ner_reg_list[end][1].startswith('I'):
                    end += 1

                ner_type = item[1].split("-")[1]

                if ner_type not in labels.keys():
                    labels[ner_type] = []

                label += "".join([item[0] for item in ner_reg_list[i:end]])
                labels[ner_type].append(label)
    return labels

def cut_text(text, lenth):
    textArr = re.findall('.{' + str(lenth) + '}', text)
    textArr.append(text[(len(textArr) * lenth):])
    return textArr
import time

def wubailong_right_company():
    company_list = []
    with open("company_right.txt",'r',encoding='utf-8') as fread:
        for line in fread:
            company = line.replace("[","").replace("]","").replace("'","").strip()
            company_list.append(company)
    for item in company_list:
        start = time.time()
        check_list = [item]
        ners = loaded_model.predict_entities(check_list)

        #from nlp_tools.utils.ner_utils import format_ner_result
        #print(format_ner_result(ners,[item]))
        #labels = extract_labels(item, ners)
        cost = time.time() - start
        item = item.split(',')
        item = item[0]
        try:
            ner_label = labels['ORG'][0]
        except Exception as e:
            ner_label = ""

        if 'ORG' in labels and labels['ORG'][0] == item:
            pass
        else:
            print("原始：%s, 提取结果：%s" % (item,ner_label))
        print("tiem cost:" +str(cost))


wubailong_right_company()


# while True:
#     #工程地点:塔山镇莒城湖东侧专业内容
#     text_input = input('sentence: ')#'本 次 招 标 项 目 的 建 设 地 点 : 双 山 岛'#input('sentence: ')
#     #text_input = text_input.split(" ")
#
#     #texts = [ '工', 程 概 况 及 招 标 范 围 1 、 建 设 地 点 : 句 容 市] #cut_text(text_input, 102)
#     text_input = list(text_input)
#     ners = loaded_model.predict([text_input])
#     #ners = loaded_model.predict([list(text_input)])
#     labels = extract_labels(text_input, ners)
#     print(labels)
#     #break
#

